import os
import numpy as np
import matplotlib.pyplot as plt
from scipy.fft import fft
from scipy.spatial.distance import cosine
import matplotlib.lines as mlines
from matplotlib.legend_handler import HandlerTuple
import pandas

plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
plt.rcParams['font.size'] = 10
plt.rcParams['axes.titlesize'] = 12
plt.rcParams['figure.titlesize'] = 14

# 创建复合图例句柄
class0_legend = (
    mlines.Line2D([], [], color='k', linewidth=3),  # 黑色边框
    mlines.Line2D([], [], color='b', linewidth=2)   # 蓝色主体
)
class1_legend = (
    mlines.Line2D([], [], color='k', linewidth=3),
    mlines.Line2D([], [], color='r', linewidth=2)
)

# 读取数据函数
def read_data_from_folder_txt(folder_path):
    data_list = []
    for file_name in os.listdir(folder_path):
        if file_name.endswith('.txt'):
            file_path = os.path.join(folder_path, file_name)
            data = np.loadtxt(file_path)
            data_list.append(data)
    return np.array(data_list)

def read_data_from_folder_csv(folder_path):
    data_list = []  # 用于存放所有样本数据（每个样本是一维数组，即电流信号）
    for file_name in os.listdir(folder_path):
        if file_name.endswith('.csv'):  # 只处理 csv 文件，不是 txt
            file_path = os.path.join(folder_path, file_name)
            try:
                # 使用 pandas 读取 csv
                df = os.read_csv(file_path)  # 默认读取所有列
                # 提取第一列数据（索引为 0），假设第一列是电流信号
                signal = df.iloc[:, 0].values  # 取第一列的所有行，转为 numpy 数组
                data_list.append(signal)
            except Exception as e:
                print(f"⚠️ 读取文件 {file_name} 时出错: {e}")
    # 将所有信号堆叠成一个 numpy 数组，shape 为 (样本数, 信号长度)
    return np.array(data_list)

# 计算时域特征
def calculate_time_domain_features(data):
    mean_val = np.mean(data, axis=1)
    std_val = np.std(data, axis=1)
    max_val = np.max(data, axis=1)
    return mean_val, std_val, max_val

# 计算频域特征
def calculate_frequency_domain_features(data):

    fft_data = np.abs(fft(data, axis=1))

    # 只取一半频率（对称性）
    n = data.shape[1]
    fft_data = fft_data[:, :n // 2]

    mean_freq = np.mean(fft_data, axis=1)
    std_freq = np.std(fft_data, axis=1)
    max_freq = np.max(fft_data, axis=1)

    return mean_freq, std_freq, max_freq

# 计算余弦相似度（样本与类均值的相似度）
def calculate_cosine_similarity_to_mean(data):

    class_mean = np.mean(data, axis=0)
    similarities = []
    for sample in data:
        cos_distance = cosine(sample, class_mean)
        cos_similarity = 1 - cos_distance
        similarities.append(cos_similarity)

    return np.array(similarities)

# 主程序
def main():
    class0_data = read_data_from_folder_txt(r"D:\A_90762\cui_python\data\test\class0")
    class1_data = read_data_from_folder_txt(r"D:\A_90762\cui_python\data\test\class1")

    print(f"Class 0 样本数: {class0_data.shape[0]}")
    print(f"Class 1 样本数: {class1_data.shape[0]}")

    # 计算时域特征
    mean0, std0, max0 = calculate_time_domain_features(class0_data)
    mean1, std1, max1 = calculate_time_domain_features(class1_data)

    # 计算频域特征
    mean_freq0, std_freq0, max_freq0 = calculate_frequency_domain_features(class0_data)
    mean_freq1, std_freq1, max_freq1 = calculate_frequency_domain_features(class1_data)

    # 计算时域余弦相似度（与类均值的相似度）
    cos_sim0 = calculate_cosine_similarity_to_mean(class0_data)
    cos_sim1 = calculate_cosine_similarity_to_mean(class1_data)

    # 计算频域余弦相似度（与类均值的相似度）
    class0_fft = np.abs(fft(class0_data, axis=1))
    class1_fft = np.abs(fft(class1_data, axis=1))
    cos_sim_freq0 = calculate_cosine_similarity_to_mean(class0_fft)
    cos_sim_freq1 = calculate_cosine_similarity_to_mean(class1_fft)

    # 计算各类特征的均值
    # 时域特征均值
    mean0_mean, mean1_mean = np.mean(mean0), np.mean(mean1)
    std0_mean, std1_mean = np.mean(std0), np.mean(std1)
    max0_mean, max1_mean = np.mean(max0), np.mean(max1)
    cos_sim0_mean, cos_sim1_mean = np.mean(cos_sim0), np.mean(cos_sim1)

    # 频域特征均值
    mean_freq0_mean, mean_freq1_mean = np.mean(mean_freq0), np.mean(mean_freq1)
    std_freq0_mean, std_freq1_mean = np.mean(std_freq0), np.mean(std_freq1)
    max_freq0_mean, max_freq1_mean = np.mean(max_freq0), np.mean(max_freq1)
    cos_sim_freq0_mean, cos_sim_freq1_mean = np.mean(cos_sim_freq0), np.mean(cos_sim_freq1)

    # 计算分类阈值（两类均值的均值）
    mean_threshold = (mean0_mean + mean1_mean) / 2
    std_threshold = (std0_mean + std1_mean) / 2
    max_threshold = (max0_mean + max1_mean) / 2
    cos_sim_threshold = (cos_sim0_mean + cos_sim1_mean) / 2

    mean_freq_threshold = (mean_freq0_mean + mean_freq1_mean) / 2
    std_freq_threshold = (std_freq0_mean + std_freq1_mean) / 2
    max_freq_threshold = (max_freq0_mean + max_freq1_mean) / 2
    cos_sim_freq_threshold = (cos_sim_freq0_mean + cos_sim_freq1_mean) / 2

    fig, axes = plt.subplots(3, 3, figsize=(20, 16))
    fig.suptitle('焊中数据时域和频域特征分析', fontsize=16)
    plt.subplots_adjust(hspace=0.4, wspace=0.3)
    # 第一行
    # 时域均值 (0,0)
    axes[0, 0].hist(mean0, bins=30, alpha=0.5, color='blue', label='Class 0')
    axes[0, 0].hist(mean1, bins=30, alpha=0.5, color='red', label='Class 1')
    axes[0, 0].axvline(mean0_mean, color='blue', linestyle='--', linewidth=2, label='Class 0 Mean')
    axes[0, 0].axvline(mean1_mean, color='red', linestyle='--', linewidth=2, label='Class 1 Mean')
    axes[0, 0].axvline(mean_threshold, color='green', linestyle='--', linewidth=2, label='Threshold')
    axes[0, 0].set_title('时域均值', pad=15)
    axes[0, 0].legend(fontsize=9)

    # 时域相似度 (余弦相似度) (0,1)
    axes[0, 1].hist(cos_sim0, bins=30, alpha=0.5, color='blue', label='Class 0')
    axes[0, 1].hist(cos_sim1, bins=30, alpha=0.5, color='red', label='Class 1')
    axes[0, 1].axvline(cos_sim0_mean, color='blue', linestyle='--', linewidth=2, label='Class 0 Mean')
    axes[0, 1].axvline(cos_sim1_mean, color='red', linestyle='--', linewidth=2, label='Class 1 Mean')
    axes[0, 1].axvline(cos_sim_threshold, color='green', linestyle='--', linewidth=2, label='Threshold')
    axes[0, 1].set_title('时域余弦相似度', pad=15)
    axes[0, 1].legend(fontsize=9)

    # 频域均值 (0,2)
    axes[0, 2].hist(mean_freq0, bins=30, alpha=0.5, color='blue', label='Class 0')
    axes[0, 2].hist(mean_freq1, bins=30, alpha=0.5, color='red', label='Class 1')
    axes[0, 2].axvline(mean_freq0_mean, color='blue', linestyle='--', linewidth=2, label='Class 0 Mean')
    axes[0, 2].axvline(mean_freq1_mean, color='red', linestyle='--', linewidth=2, label='Class 1 Mean')
    axes[0, 2].axvline(mean_freq_threshold, color='green', linestyle='--', linewidth=2, label='Threshold')
    axes[0, 2].set_title('频域均值', pad=15)
    axes[0, 2].legend(fontsize=9)

    # 第二行
    # 时域标准差 (1,0)
    axes[1, 0].hist(std0, bins=30, alpha=0.5, color='blue', label='Class 0')
    axes[1, 0].hist(std1, bins=30, alpha=0.5, color='red', label='Class 1')
    axes[1, 0].axvline(std0_mean, color='blue', linestyle='--', linewidth=2, label='Class 0 Mean')
    axes[1, 0].axvline(std1_mean, color='red', linestyle='--', linewidth=2, label='Class 1 Mean')
    axes[1, 0].axvline(std_threshold, color='green', linestyle='--', linewidth=2, label='Threshold')
    axes[1, 0].set_title('时域标准差', pad=15)
    axes[1, 0].legend(fontsize=9)

    # 原始信号 (1,1)
    time_axis = np.arange(500)

    for signal in class0_data:
        axes[1, 1].plot(time_axis, signal, 'b-', alpha=0.1)
    for signal in class1_data:
        axes[1, 1].plot(time_axis, signal, 'r-', alpha=0.1)
    class0_mean = np.mean(class0_data, axis=0)
    class1_mean = np.mean(class1_data, axis=0)
    axes[1, 1].plot(time_axis, class0_mean, 'k-', linewidth=3)  # 黑色边框
    axes[1, 1].plot(time_axis, class1_mean, 'k-', linewidth=3)
    axes[1, 1].plot(time_axis, class0_mean, 'b-', linewidth=2)
    axes[1, 1].plot(time_axis, class1_mean, 'r-', linewidth=2)

    axes[1, 1].set_title('原始信号', pad=15)
    axes[1, 1].set_xlabel('时间点', fontsize=10)
    axes[1, 1].set_ylabel('幅值', fontsize=10)
    axes[1, 1].legend(
        handles=[class0_legend, class1_legend],
        labels=['Class 0 Mean', 'Class 1 Mean'],
        handler_map={tuple: HandlerTuple(ndivide=1)},
        loc='upper right'
    )

    # 频域标准差 (1,2)
    axes[1, 2].hist(std_freq0, bins=30, alpha=0.5, color='blue', label='Class 0')
    axes[1, 2].hist(std_freq1, bins=30, alpha=0.5, color='red', label='Class 1')
    axes[1, 2].axvline(std_freq0_mean, color='blue', linestyle='--', linewidth=2, label='Class 0 Mean')
    axes[1, 2].axvline(std_freq1_mean, color='red', linestyle='--', linewidth=2, label='Class 1 Mean')
    axes[1, 2].axvline(std_freq_threshold, color='green', linestyle='--', linewidth=2, label='Threshold')
    axes[1, 2].set_title('频域标准差', pad=15)
    axes[1, 2].legend(fontsize=9)

    # 第三行
    # 时域最大值 (2,0)
    axes[2, 0].hist(max0, bins=30, alpha=0.5, color='blue', label='Class 0')
    axes[2, 0].hist(max1, bins=30, alpha=0.5, color='red', label='Class 1')
    axes[2, 0].axvline(max0_mean, color='blue', linestyle='--', linewidth=2, label='Class 0 Mean')
    axes[2, 0].axvline(max1_mean, color='red', linestyle='--', linewidth=2, label='Class 1 Mean')
    axes[2, 0].axvline(max_threshold, color='green', linestyle='--', linewidth=2, label='Threshold')
    axes[2, 0].set_title('时域最大值', pad=15)
    axes[2, 0].legend(fontsize=9)

    # 频域相似度 (余弦相似度) (2,1)
    axes[2, 1].hist(cos_sim_freq0, bins=30, alpha=0.5, color='blue', label='Class 0')
    axes[2, 1].hist(cos_sim_freq1, bins=30, alpha=0.5, color='red', label='Class 1')
    axes[2, 1].axvline(cos_sim_freq0_mean, color='blue', linestyle='--', linewidth=2, label='Class 0 Mean')
    axes[2, 1].axvline(cos_sim_freq1_mean, color='red', linestyle='--', linewidth=2, label='Class 1 Mean')
    axes[2, 1].axvline(cos_sim_freq_threshold, color='green', linestyle='--', linewidth=2, label='Threshold')
    axes[2, 1].set_title('频域余弦相似度', pad=15)
    axes[2, 1].legend(fontsize=9)

    # 频域最大值 (2,2)
    axes[2, 2].hist(max_freq0, bins=30, alpha=0.5, color='blue', label='Class 0')
    axes[2, 2].hist(max_freq1, bins=30, alpha=0.5, color='red', label='Class 1')
    axes[2, 2].axvline(max_freq0_mean, color='blue', linestyle='--', linewidth=2, label='Class 0 Mean')
    axes[2, 2].axvline(max_freq1_mean, color='red', linestyle='--', linewidth=2, label='Class 1 Mean')
    axes[2, 2].axvline(max_freq_threshold, color='green', linestyle='--', linewidth=2, label='Threshold')
    axes[2, 2].set_title('频域最大值', pad=15)
    axes[2, 2].legend(fontsize=9)

    plt.tight_layout()
    plt.subplots_adjust(top=0.93)

    plt.savefig('welding_features_analysis.png', dpi=300, bbox_inches='tight')
    print("图像已保存为 'welding_features_analysis.png'")
    plt.show()

    print("\n特征统计信息:")
    print("时域均值 - Class 0: {:.4f}, Class 1: {:.4f}, 阈值: {:.4f}".format(mean0_mean, mean1_mean, mean_threshold))
    print("时域标准差 - Class 0: {:.4f}, Class 1: {:.4f}, 阈值: {:.4f}".format(std0_mean, std1_mean, std_threshold))
    print("时域最大值 - Class 0: {:.4f}, Class 1: {:.4f}, 阈值: {:.4f}".format(max0_mean, max1_mean, max_threshold))
    print("时域余弦相似度 - Class 0: {:.4f}, Class 1: {:.4f}, 阈值: {:.4f}".format(cos_sim0_mean, cos_sim1_mean,
                                                                                   cos_sim_threshold))
    print("频域均值 - Class 0: {:.4f}, Class 1: {:.4f}, 阈值: {:.4f}".format(mean_freq0_mean, mean_freq1_mean,
                                                                             mean_freq_threshold))
    print("频域标准差 - Class 0: {:.4f}, Class 1: {:.4f}, 阈值: {:.4f}".format(std_freq0_mean, std_freq1_mean,
                                                                               std_freq_threshold))
    print("频域最大值 - Class 0: {:.4f}, Class 1: {:.4f}, 阈值: {:.4f}".format(max_freq0_mean, max_freq1_mean,
                                                                               max_freq_threshold))
    print(
        "频域余弦相似度 - Class 0: {:.4f}, Class 1: {:.4f}, 阈值: {:.4f}".format(cos_sim_freq0_mean, cos_sim_freq1_mean,
                                                                                 cos_sim_freq_threshold))

if __name__ == "__main__":
    # 请确保路径正确
    class0_train_path = r"D:\A_90762\cui_python\data\train\class0"
    class1_train_path = r"D:\A_90762\cui_python\data\train\class1"

    main()

